import cv2
import gi
gi.require_version('Gst', '1.0')
gi.require_version('GstRtspServer', '1.0')
from gi.repository import Gst, GstRtspServer, GLib, GObject
WIDTH = 1920
HEIGHT = 1080
class MyContext():
    def __init__(self):
        self.white = False
        self.timestamp = 0

Gst.init(None)

# 打开摄像头
cap = cv2.VideoCapture(0)
# 设置摄像头参数
cap.set(cv2.CAP_PROP_FRAME_WIDTH, WIDTH)  # 设置宽度
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, HEIGHT)  # 设置高度
# cap.set(cv2.CAP_PROP_FPS, 30)  # 设置帧率

# 创建factory
factory = GstRtspServer.RTSPMediaFactory()
factory.set_launch("( appsrc name=videosrc is-live=true ! videoconvert ! nvvidconv ! video/x-raw(memory:NVMM), format=I420 ! nvv4l2h264enc maxperf-enable=1 control-rate=0 bitrate=4000000 preset-level=2 profile=2 ! rtph264pay name=pay0 pt=96 )")
factory.set_shared(True)

def media_configure(factory, media):
    print("media_configure")
    element = media.get_element()

    # 获取appsrc元素
    appsrc = element.get_by_name("videosrc")
    appsrc.set_property("stream-type", 0)
    appsrc.set_property("format", Gst.Format.TIME)

    caps = Gst.Caps.from_string("video/x-raw,format=BGR,width={},height={},framerate=30/1,pixel-aspect-ratio=1/1".format(WIDTH, HEIGHT))
    appsrc.set_property("caps", caps)

    ctx = MyContext()
    ctx.white = False
    ctx.timestamp = 0
    # media.set_user_data(ctx, MyContext)

    # 连接"need-data"信号和回调函数
    appsrc.connect("need-data", push_frame, ctx)

# 连接信号和回调函数
factory.connect("media-configure", media_configure)

# 创建RTSP服务器
server = GstRtspServer.RTSPServer()
server.set_service("8554")
mounts = server.get_mount_points()

mounts.add_factory("/test", factory)
server.attach(None)

# 定义回调函数，用于向appsrc元素推送视频帧
def push_frame(appsrc, length, ctx):
    ret, frame = cap.read()
    # frame = cv2.resize(frame,(1920,1080),interpolation=cv2.INTER_NEAREST)
    if ret:
        print(frame.shape)
        # 将OpenCV的图像数据转换为GStreamer需要的格式
        data = frame.tobytes()
        buffer = Gst.Buffer.new_allocate(None, len(data), None)
        buffer.fill(0, data)
        buffer.pts = ctx.timestamp
        buffer.duration = Gst.util_uint64_scale_int(1, Gst.SECOND, 30)
        ctx.timestamp+=buffer.duration

        # 向appsrc元素推送视频帧
        result = appsrc.emit("push-buffer", buffer)

        # 判断推送结果
        # if result == Gst.FlowReturn.OK:
        #     print("Buffer pushed successfully")
        # else:
        #     print("Failed to push buffer")

    return True

# 进入主循环
loop = GLib.MainLoop()
loop.run()

# 停止GStreamer管道
# pipeline.set_state(Gst.State.NULL)

# 释放摄像头
cap.release()
# 释放挂载点
mounts.remove_factory("/test")
